import torch
import torch.nn as nn


class MF(nn.Module):
    def __init__(self, hyper_params):
        super(MF, self).__init__()
        self.hyper_parmas = hyper_params

        self.item_embedding = nn.Embedding(
            num_embeddings=hyper_params['item_cnt'],
            embedding_dim=hyper_params['embed_dim']
        )
        self.user_embedding = nn.Embedding(
            num_embeddings=hyper_params['user_cnt'],
            embedding_dim=hyper_params['embed_dim']
        )

    def forward(self, item_id, user_id):
        item_embed = self.item_embedding(item_id)
        user_embed = self.user_embedding(user_id)
        result = torch.sum(item_embed * user_embed, dim=-1)

        return result

    def get_user_embed(self, user_id):
        user_embed = self.user_embedding(user_id)
        return user_embed

    def get_item_embed(self, item_id):
        item_embed = self.item_embedding(item_id)
        return item_embed


class PairwiseLoss:
    def __init__(self, hyper_params, mf):
        self.hyper_params = hyper_params
        self.MF = mf

    def get_loss(self, pos_item, neg_item, user_id):
        pos_val = self.MF(pos_item, user_id)
        neg_val = self.MF(neg_item, user_id)

        result = pos_val - neg_val
        loss = -torch.mean(torch.log(torch.sigmoid(result)+1e-7))

        return loss
